'''
Author: ironlionwgHH huanghao_bnu@163.com
Date: 2023-05-24 15:21:39
LastEditors: HuangHao ironlionwg@outlook.com
LastEditTime: 2024-09-05 17:44:16
FilePath: /asteroid_orbit/_cartesian_oep_process.py
Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
'''
'''
 _   _                                     _   _                   _          ____               _
| | | |  _   _    __ _   _ __     __ _    | | | |   __ _    ___   ( )___     / ___|   ___     __| |   ___ 
| |_| | | | | |  / _` | | '_ \   / _` |   | |_| |  / _` |  / _ \  |// __|   | |      / _ \   / _` |  / _ \
|  _  | | |_| | | (_| | | | | | | (_| |   |  _  | | (_| | | (_) |   \__ \   | |___  | (_) | | (_| | |  __/
|_| |_|  \__,_|  \__,_| |_| |_|  \__, |   |_| |_|  \__,_|  \___/    |___/    \____|  \___/   \__,_|  \___|
                                  |___/
'''
import os
import pandas as pd
import spiceypy as sp
from matplotlib import pyplot as plt
import numpy as np

AU_km = 149597870.700

'''
Class    :
            cartesian_oep_file
Description:
            Process cartesian oep file contains cartesian coordinate in ECLM
Edition  :  黄皓/2023.05.24/
Log      :
__init__ :
            cartesian_oep_file_path     str     Path    | e.g. "./2016HO3"
            cartesian_oep_file_name     str     Name    | e.g. "cartesian_eph.oep"
            eph_num                     int     Hwo many useful eph contain in this file
Method   :
            
Warning  :
            
'''
class cartesian_oep_file():
    def __init__(self, cartesian_oep_file_path: str, cartesian_oep_file_name: str, eph_num: int):
        self.cartesian_oep_file_path = cartesian_oep_file_path
        self.cartesian_oep_file_name = cartesian_oep_file_name
        self.cartesian_oep_file_path_name = os.path.join(self.cartesian_oep_file_path, self.cartesian_oep_file_name)
        self.object_name = cartesian_oep_file_name.split("_")[0]
        # 提取信息
        self.car_eph_dataframe = pd.read_csv(self.cartesian_oep_file_path_name, sep = "\\s+", header=None)
        self.car_eph_dataframe_eph = self.car_eph_dataframe[-eph_num:]
    
    '''
    Function :
                read_the_eph
    Description :
                Read the cartesian eph and return a DataFrame type sorted by the date
    Edition  :  Huang Hao/2023.05.24/
    Log      :
    Input    :
    Return   :
                car_eph_dataframe_sort      DataFrame   a DataFrame type sorted by the date
    func used:
    Warning  :
                
    '''
    def read_the_eph(self):
        car_eph_dataframe_sort = self.car_eph_dataframe_eph.sort_values(by=[0],ascending=[True])
        
        return car_eph_dataframe_sort
    
    def draw_diff_horizon(self, begin_year, title, fig_path, fig_name,obj_id):
        car_eph_dataframe_sort = self.read_the_eph()
        time = np.array(car_eph_dataframe_sort[0])
        
        
        begin_jd = car_eph_dataframe_sort.iloc[0][0] + 2400000.5
        end_jd   = car_eph_dataframe_sort.iloc[-1][0] + 2400000.5

        for i in time:
            tdb = sp.str2et("JD" + str(i+2400000.5) + " TDT")   #Orbfit里最后一列是TDT时间，用这个函数转换成TDB时间（单位：秒）
            spk_information = sp.spkezr(str(obj_id), tdb, 'ECLIPJ2000', 'None', '10')[0].reshape(1,6)
            if i == begin_jd-2400000.5:
                spk_data = spk_information
            else:
                spk_data = np.vstack((spk_data, spk_information))
        r_orbfit_spice = np.linalg.norm( (np.array(car_eph_dataframe_sort)[:,1:4]*AU_km - spk_data[:,:3]) ,axis=1)
        v_orbfit_spice = np.linalg.norm( (np.array(car_eph_dataframe_sort)[:,4:7]*AU_km/86400 - spk_data[:,3:6]) ,axis=1)
        
        r_spice = np.linalg.norm(spk_data[:,:3],axis=1); r_orbfit =np.linalg.norm(np.array(car_eph_dataframe_sort)[:,1:4],axis=1)
        v_spice = np.linalg.norm(spk_data[:,3:6],axis=1); v_orbfit =np.linalg.norm(np.array(car_eph_dataframe_sort)[:,4:7],axis=1)
        
        
        x_spice = spk_data.T[0]; x_orbfit = car_eph_dataframe_sort[1]
        y_spice = spk_data.T[1]; y_orbfit = car_eph_dataframe_sort[2]
        z_spice = spk_data.T[2]; z_orbfit = car_eph_dataframe_sort[3]
        
        vx_spice = spk_data.T[3]; vx_orbfit = car_eph_dataframe_sort[4]
        vy_spice = spk_data.T[4]; vy_orbfit = car_eph_dataframe_sort[5]
        vz_spice = spk_data.T[5]; vz_orbfit = car_eph_dataframe_sort[6]

        # ++++++++++++++++++++++++++++++++++++++++ Draw the result ++++++++++++++++++++++++++++++++++++++++
        time = (time-time[0])/365+begin_year
        # ========== 位置差异 ==========
        if not os.path.exists(fig_path):
            os.makedirs(fig_path)
        legend_size = 10; title_size = 20; label_size = 15; tick_size = 15
        font_legend = {'family' : 'Times New Roman',
                        'weight' : 'heavy',
                        'size'   : legend_size,}
        font_title = {'family' : 'Times New Roman',
                    'weight' : 'heavy',
                    'size'   : title_size,}
        font_label = {'family' : 'Times New Roman',
                    'weight' : 'heavy',
                    'size'   : label_size,}
        font_tick = {'family' : 'Times New Roman',
                        'weight' : 'heavy',
                        'size'   : tick_size,}
        fig, ax = plt.subplots(2, 1, figsize=(8, 3),sharex=True)
        plt.subplots_adjust(hspace=0.1)
        ax[0].plot(time, -x_orbfit*AU_km+x_spice, label = "x"); rms_x = np.sqrt(np.mean(np.square(-x_orbfit*AU_km+x_spice)))
        ax[0].plot(time, -y_orbfit*AU_km+y_spice, label = "y"); rms_y = np.sqrt(np.mean(np.square(-y_orbfit*AU_km+y_spice)))
        ax[0].plot(time, -z_orbfit*AU_km+z_spice, label = "z"); rms_z = np.sqrt(np.mean(np.square(-z_orbfit*AU_km+z_spice)))
        ax[0].plot(time, r_orbfit_spice, label = "r")         ; rms_r = np.sqrt(np.mean(np.square(r_orbfit_spice)))
        ax[0].set_ylabel(r"$\Delta / km$", font_label)
        tex_box_props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)
        # ax[0].text(0.83, 0.95, "RMS_x = %.2f\nRMS_y = %.2f\nRMS_z = %.2f\nRMS_r = %.2f"%(rms_x,rms_y,rms_z,rms_r), 
        #            transform=ax[0].transAxes, fontdict=font_legend,
        #            verticalalignment='top', bbox=tex_box_props)
        ax[0].set_title(title, fontdict = font_title)

        # ========== 速度差异 ==========
        ax[1].plot(time, -(vx_orbfit*AU_km/86400 - vx_spice)*1e5, label = r"$v_{X}$"); rms_vx = np.sqrt(np.mean(np.square(-(vx_orbfit*AU_km/86400 - vx_spice)*1e5)))
        ax[1].plot(time, -(vy_orbfit*AU_km/86400 - vy_spice)*1e5, label = r"$v_{Y}$"); rms_vy = np.sqrt(np.mean(np.square(-(vy_orbfit*AU_km/86400 - vy_spice)*1e5)))
        ax[1].plot(time, -(vz_orbfit*AU_km/86400 - vz_spice)*1e5, label = r"$v_{Z}$"); rms_vz = np.sqrt(np.mean(np.square(-(vz_orbfit*AU_km/86400 - vz_spice)*1e5)))
        ax[1].plot(time, v_orbfit_spice*1e5, label = r"$v$")                         ; rms_v = np.sqrt(np.mean(np.square(v_orbfit_spice*1e5)))
        # ax[1].text(0.83, 0.95, "RMS_vx = %.2f\nRMS_vy = %.2f\nRMS_vz = %.2f\nRMS_v  = %.2f"%(rms_vx,rms_vy,rms_vz,rms_v), 
        #            transform=ax[1].transAxes, fontdict=font_legend,
        #            verticalalignment='top', bbox=tex_box_props)
        ax[1].tick_params(direction = "in", labelsize = tick_size)
        ax[1].set_xlabel(r"Time / year", fontdict = font_label)
        ax[1].set_ylabel(r"$\Delta / cm\cdot{s}^{-1}$", font_label)
        
        for sub_ax in ax:
            sub_ax.legend(prop = font_legend,ncol=3,frameon=False)
            sub_ax.tick_params(direction = "in", labelsize = tick_size)
        print("RMS relative to Horizon:")
        print("RMS_x  =%10.5f RMS_y  =%10.5f RMS_z  =%10.5f RMS_r =%10.5f"%(rms_x,rms_y,rms_z,rms_r))
        print("RMS_vx =%10.5f RMS_vy =%10.5f RMS_vz =%10.5f RMS_v =%10.5f"%(rms_vx,rms_vy,rms_vz,rms_v))
        plt.savefig("%s.png"%(os.path.join(fig_path, fig_name)), format = "png", bbox_inches = "tight", dpi = 1000)
        plt.savefig("%s.pdf"%(os.path.join(fig_path, fig_name)), format = "pdf", bbox_inches = "tight", dpi = 1000)
        plt.savefig("%s.svg"%(os.path.join(fig_path, fig_name)), format = "svg", bbox_inches = "tight", dpi = 1000)
        plt.close()

if __name__ == "__main__":
    test_file = cartesian_oep_file("./2016HO3","cartesian_eph.oep",18629)
    sp.furnsh("spice_kernel/2469219.bsp")
    sp.furnsh("spice_kernel/de440.bsp")
    sp.furnsh("spice_kernel/naif0012.tls.pc")
    car_eph_dataframe_sort = test_file.read_the_eph()
    test_file.draw_diff_horizon(begin_year=2000, title="test", fig_path="./", fig_name="test_draw")